#include "flashros/inference/input_pipeline.h"

using namespace flashros;
using namespace inference;

void RosInput::Init(std::string camera_topic_name, int max_wait_secs,
                    ros::NodeHandle& nh) {
  // ros subscriber callback function
  auto callback_fn = [&](const sensor_msgs::ImageConstPtr& msg) {
    try {
      lock_.lock();
      frames_.clear();
      frames_.push_back(cv_bridge::toCvShare(msg, "bgr8")->image);
      SPDLOG_INFO("Add one frame with shape ({}, {})", frames_.back().rows,
                  frames_.back().cols);
      lock_.unlock();
    } catch (cv_bridge::Exception& e) {
      SPDLOG_ERROR("Could not convert from '{}' to 'bgr8'.",
                   msg->encoding.c_str());
    }
  };

  max_wait_secs_ = max_wait_secs;
  frames_ = std::vector<cv::Mat>();
  running_ = true;

  // start a new thread to run subscriber
  thread_ = std::thread([&]() {
    image_transport::ImageTransport it(nh);
    image_transport::Subscriber sub =
        it.subscribe(camera_topic_name, 1, callback_fn);
    SPDLOG_INFO("Ros Input thread started");
    while (running_) {
      ros::spinOnce();
    }
  });

  std::this_thread::sleep_for(std::chrono::seconds(10));
  SPDLOG_INFO("Ros input pipeline init successfully with topic name '{}'",
              camera_topic_name);
  std::this_thread::sleep_for(std::chrono::seconds(10));
}

bool RosInput::Read(cv::Mat& frame) {
  float secs = 0;
  while (secs < max_wait_secs_) {
    lock_.lock();
    if (frames_.size() > 0) {
      frame = frames_.back();
      frames_.pop_back();
      SPDLOG_INFO("Ros Input receive one frame");
      lock_.unlock();
      return true;
    }
    lock_.unlock();
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    secs += 0.05;
    SPDLOG_INFO("Ros Input failed to get one image for {} secs.", secs);
  }
  return false;
}